#include <mcp/client/client.h>
#include <mcp/transport/transport.h>
#include <mcp/json/value.h>
#include <mcp/error.h>
#include <mcp/types.h>
#include <mcp/forward.h>

#include <iostream>
#include <sstream>
#include <chrono>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <future>
#include <memory>
#include <unordered_map>

namespace mcp {
namespace client {

// Client implementation
struct Client::Impl {
    std::unique_ptr<Transport> transport;
    ClientConfig config;
    ClientCapabilities capabilities;
    bool connected = false;
    std::string server_protocol_version;
    json::Value server_info;
    json::Value server_capabilities;

    // Message handling
    std::mutex message_mutex;
    std::unordered_map<std::string, std::promise<json::Value>> pending_requests;
    std::function<void(const std::string&, const json::Value&)> notification_handler;
    std::function<void(const Error&)> error_handler;

    // Request ID management
    std::atomic<int> next_request_id{1};

    // Threading
    std::thread receiver_thread;
    std::atomic<bool> running{false};

    // Tool output schema cache
    std::unordered_map<std::string, json::Value> tool_output_schemas;
    mutable std::mutex schema_mutex;

    Impl(std::unique_ptr<Transport> trans, const ClientConfig& cfg)
        : transport(std::move(trans)), config(cfg) {}

    std::string generate_request_id() {
        return std::to_string(next_request_id.fetch_add(1));
    }

    void handle_message(const json::Value& message);
    void handle_request(const json::Value& request);
    void handle_response(const json::Value& response);
    void handle_notification(const std::string& method, const json::Value& params);

    void receiver_loop();
};

Client::Client(std::unique_ptr<Transport> transport, const ClientConfig& config)
    : impl_(std::make_unique<Impl>(std::move(transport), config)) {}

Client::~Client() {
    disconnect();
}

bool Client::connect() {
    if (impl_->connected) {
        return true;
    }

    if (!impl_->transport->is_connected()) {
        if (!impl_->transport->connect()) {
            if (impl_->error_handler) {
                impl_->error_handler(Error::connection_error("Failed to connect to transport"));
            }
            return false;
        }
    }

    // Send initialize request
    auto init_params = json::Value::object({
        {"protocolVersion", "2025-06-18"},
        {"capabilities", json::Value::object({
            {"tools", impl_->capabilities.tools},
            {"resources", impl_->capabilities.resources},
            {"prompts", impl_->capabilities.prompts},
            {"logging", impl_->capabilities.logging},
            {"progress", impl_->capabilities.progress},
            {"sampling", impl_->capabilities.sampling},
            {"roots", true} // Always support roots
        })},
        {"clientInfo", json::Value::object({
            {"name", impl_->config.name},
            {"version", impl_->config.version}
        })}
    });

    auto request = json::Value::object({
        {"method", "initialize"},
        {"params", std::move(init_params)}
    });

    // Send initialize request
    std::promise<json::Value> init_promise;
    auto init_future = init_promise.get_future();

    {
        std::lock_guard<std::mutex> lock(impl_->message_mutex);
        auto request_id = impl_->generate_request_id();
        request["id"] = request_id;
        impl_->pending_requests[request_id] = std::move(init_promise);

        if (!impl_->transport->send_message(request.serialize())) {
            if (impl_->error_handler) {
                impl_->error_handler(Error::connection_error("Failed to send initialize request"));
            }
            return false;
        }
    }

    // Wait for initialize response
    if (init_future.wait_for(std::chrono::seconds(impl_->config.timeout_seconds)) !=
        std::future_status::ready) {
        if (impl_->error_handler) {
            impl_->error_handler(Error::timeout("Initialize request timed out"));
        }
        return false;
    }

    auto response = init_future.get();
    if (!response.is_object() || !response.contains("result")) {
        if (impl_->error_handler) {
            impl_->error_handler(Error::protocol_error("Invalid initialize response"));
        }
        return false;
    }

    auto result = response["result"];
    impl_->server_protocol_version = result.get("protocolVersion", "").as_string();
    impl_->server_info = result["serverInfo"];
    impl_->server_capabilities = result["capabilities"];

    // Send initialized notification
    auto initialized_notification = json::Value::object({
        {"method", "initialized"},
        {"params", json::Value::object({})}
    });

    if (!impl_->transport->send_message(initialized_notification.serialize())) {
        if (impl_->error_handler) {
            impl_->error_handler(Error::connection_error("Failed to send initialized notification"));
        }
        return false;
    }

    // Start receiver thread
    impl_->running = true;
    impl_->receiver_thread = std::thread(&Impl::receiver_loop, impl_.get());

    impl_->connected = true;
    return true;
}

void Client::disconnect() {
    if (!impl_->connected) {
        return;
    }

    impl_->running = false;
    impl_->connected = false;

    if (impl_->receiver_thread.joinable()) {
        impl_->receiver_thread.join();
    }

    impl_->transport->disconnect();

    // Clear pending requests
    std::lock_guard<std::mutex> lock(impl_->message_mutex);
    for (auto& [id, promise] : impl_->pending_requests) {
        promise.set_value(Error::disconnected("Connection closed").to_json());
    }
    impl_->pending_requests.clear();
}

bool Client::is_connected() const {
    return impl_->connected && impl_->transport->is_connected();
}

Result<json::Value> Client::get_server_info() {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }
    return impl_->server_info;
}

Result<json::Value> Client::get_capabilities() {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }
    return impl_->server_capabilities;
}

Result<std::string> Client::get_protocol_version() {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }
    return impl_->server_protocol_version;
}

Result<std::vector<types::Tool>> Client::list_tools() {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }

    auto request = json::Value::object({
        {"method", "tools/list"},
        {"params", json::Value::object({})}
    });

    auto response = send_request(request);
    if (!response) {
        return response.error();
    }

    auto result = response.value().get("result", json::Value::object({}));
    auto tools_array = result.get("tools", json::Value::array({}));

    std::vector<types::Tool> tools;
    for (const auto& tool_json : tools_array) {
        types::Tool tool;
        tool.name = tool_json.get("name", "").as_string();
        tool.description = tool_json.get("description", "").as_string();
        tool.input_schema = tool_json.get("inputSchema", json::Value::object({}));

        // Cache output schema
        if (tool_json.contains("outputSchema")) {
            std::lock_guard<std::mutex> lock(impl_->schema_mutex);
            impl_->tool_output_schemas[tool.name] = tool_json["outputSchema"];
        }

        tools.push_back(std::move(tool));
    }

    return tools;
}

Result<json::Value> Client::call_tool(const std::string& name, const json::Value& arguments) {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }

    auto params = json::Value::object({
        {"name", name},
        {"arguments", arguments}
    });

    auto request = json::Value::object({
        {"method", "tools/call"},
        {"params", std::move(params)}
    });

    auto response = send_request(request);
    if (!response) {
        return response.error();
    }

    auto result = response.value().get("result", json::Value::object({}));

    // Validate structured content if tool has output schema
    {
        std::lock_guard<std::mutex> lock(impl_->schema_mutex);
        if (impl_->tool_output_schemas.contains(name)) {
            auto output_schema = impl_->tool_output_schemas[name];
            if (result.contains("content") && result["content"].is_array()) {
                // Basic validation - in real implementation would use proper schema validation
                auto content = result["content"];
                for (const auto& item : content) {
                    if (!item.is_object() || !item.contains("type")) {
                        return Error::validation_error("Invalid structured content format");
                    }
                }
            }
        }
    }

    return result;
}

Result<std::vector<types::Resource>> Client::list_resources() {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }

    auto request = json::Value::object({
        {"method", "resources/list"},
        {"params", json::Value::object({})}
    });

    auto response = send_request(request);
    if (!response) {
        return response.error();
    }

    auto result = response.value().get("result", json::Value::object({}));
    auto resources_array = result.get("resources", json::Value::array({}));

    std::vector<types::Resource> resources;
    for (const auto& resource_json : resources_array) {
        types::Resource resource;
        resource.uri = resource_json.get("uri", "").as_string();
        resource.name = resource_json.get("name", "").as_string();
        resource.description = resource_json.get("description", "").as_string();
        resource.mime_type = resource_json.get("mimeType", "").as_string();
        resources.push_back(std::move(resource));
    }

    return resources;
}

Result<std::vector<types::ResourceTemplate>> Client::list_resource_templates() {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }

    auto request = json::Value::object({
        {"method", "resources/templates/list"},
        {"params", json::Value::object({})}
    });

    auto response = send_request(request);
    if (!response) {
        return response.error();
    }

    auto result = response.value().get("result", json::Value::object({}));
    auto templates_array = result.get("resourceTemplates", json::Value::array({}));

    std::vector<types::ResourceTemplate> templates;
    for (const auto& template_json : templates_array) {
        types::ResourceTemplate templ;
        templ.uri_template = template_json.get("uriTemplate", "").as_string();
        templ.name = template_json.get("name", "").as_string();
        templ.description = template_json.get("description", "").as_string();
        templ.input_schema = template_json.get("inputSchema", json::Value::object({}));
        templates.push_back(std::move(templ));
    }

    return templates;
}

Result<types::ResourceContent> Client::read_resource(const std::string& uri) {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }

    auto params = json::Value::object({
        {"uri", uri}
    });

    auto request = json::Value::object({
        {"method", "resources/read"},
        {"params", std::move(params)}
    });

    auto response = send_request(request);
    if (!response) {
        return response.error();
    }

    auto result = response.value().get("result", json::Value::object({}));
    auto contents_array = result.get("contents", json::Value::array({}));

    if (contents_array.empty()) {
        return Error::protocol_error("No content returned for resource");
    }

    auto content_json = contents_array[0];
    types::ResourceContent content;
    content.uri = content_json.get("uri", "").as_string();
    content.mime_type = content_json.get("mimeType", "").as_string();

    if (content_json.contains("text")) {
        content.content = content_json["text"].as_string();
    } else if (content_json.contains("blob")) {
        content.content = content_json["blob"].as_string(); // Base64 encoded
    } else {
        return Error::protocol_error("Invalid resource content format");
    }

    return content;
}

Result<std::vector<types::Prompt>> Client::list_prompts() {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }

    auto request = json::Value::object({
        {"method", "prompts/list"},
        {"params", json::Value::object({})}
    });

    auto response = send_request(request);
    if (!response) {
        return response.error();
    }

    auto result = response.value().get("result", json::Value::object({}));
    auto prompts_array = result.get("prompts", json::Value::array({}));

    std::vector<types::Prompt> prompts;
    for (const auto& prompt_json : prompts_array) {
        types::Prompt prompt;
        prompt.name = prompt_json.get("name", "").as_string();
        prompt.description = prompt_json.get("description", "").as_string();
        prompt.arguments = prompt_json.get("arguments", json::Value::array({}));
        prompts.push_back(std::move(prompt));
    }

    return prompts;
}

Result<std::vector<types::PromptMessage>> Client::call_prompt(const std::string& name,
                                                             const json::Value& arguments) {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }

    auto params = json::Value::object({
        {"name", name},
        {"arguments", arguments}
    });

    auto request = json::Value::object({
        {"method", "prompts/get"},
        {"params", std::move(params)}
    });

    auto response = send_request(request);
    if (!response) {
        return response.error();
    }

    auto result = response.value().get("result", json::Value::object({}));
    auto messages_array = result.get("messages", json::Value::array({}));

    std::vector<types::PromptMessage> messages;
    for (const auto& message_json : messages_array) {
        types::PromptMessage message;
        message.role = message_json.get("role", "").as_string();

        if (message_json.contains("content")) {
            auto content = message_json["content"];
            if (content.is_object()) {
                message.content = content.get("text", "").as_string();
            } else {
                message.content = content.as_string();
            }
        }

        messages.push_back(std::move(message));
    }

    return messages;
}

bool Client::send_log(types::LogLevel level, const std::string& message,
                     const std::optional<std::string>& logger_name) {
    if (!impl_->connected) {
        return false;
    }

    std::string level_str;
    switch (level) {
        case types::LogLevel::Debug: level_str = "debug"; break;
        case types::LogLevel::Info: level_str = "info"; break;
        case types::LogLevel::Warning: level_str = "warning"; break;
        case types::LogLevel::Error: level_str = "error"; break;
        default: level_str = "info"; break;
    }

    auto params = json::Value::object({
        {"level", level_str},
        {"data", message}
    });

    if (logger_name) {
        params["logger"] = *logger_name;
    }

    auto notification = json::Value::object({
        {"method", "logging/message"},
        {"params", std::move(params)}
    });

    return impl_->transport->send_message(notification.serialize());
}

bool Client::send_progress(const types::ProgressToken& token, double progress,
                          std::optional<double> total, std::optional<std::string> message) {
    if (!impl_->connected) {
        return false;
    }

    auto params = json::Value::object({
        {"progressToken", std::visit([](auto&& arg) -> json::Value {
            using T = std::decay_t<decltype(arg)>;
            if constexpr (std::is_same_v<T, std::string>) {
                return arg;
            } else {
                return static_cast<int64_t>(arg);
            }
        }, token)},
        {"progress", progress}
    });

    if (total) {
        params["total"] = *total;
    }

    if (message) {
        params["message"] = *message;
    }

    auto notification = json::Value::object({
        {"method", "notifications/progress"},
        {"params", std::move(params)}
    });

    return impl_->transport->send_message(notification.serialize());
}

Result<types::SamplingResult> Client::create_message(
    const std::vector<types::SamplingMessage>& messages,
    const types::SamplingOptions& options) {
    if (!impl_->connected) {
        return Error::not_connected("Not connected to server");
    }

    auto messages_array = json::Value::array({});
    for (const auto& msg : messages) {
        auto msg_obj = json::Value::object({
            {"role", msg.role},
            {"content", msg.content}
        });

        if (msg.name) {
            msg_obj["name"] = *msg.name;
        }

        messages_array.push_back(std::move(msg_obj));
    }

    auto params = json::Value::object({
        {"messages", std::move(messages_array)},
        {"systemPrompt", options.system_prompt.value_or("")},
        {"temperature", options.temperature},
        {"maxTokens", options.max_tokens}
    });

    if (options.stop_sequences) {
        auto stop_array = json::Value::array({});
        for (const auto& stop : *options.stop_sequences) {
            stop_array.push_back(stop);
        }
        params["stop"] = std::move(stop_array);
    }

    auto request = json::Value::object({
        {"method", "sampling/createMessage"},
        {"params", std::move(params)}
    });

    auto response = send_request(request);
    if (!response) {
        return response.error();
    }

    auto result = response.value().get("result", json::Value::object({}));

    types::SamplingResult sampling_result;
    sampling_result.content = result.get("content", "").as_string();
    sampling_result.model = result.get("model", "").as_string();
    sampling_result.stop_reason = result.get("stopReason", "").as_string();

    return sampling_result;
}

void Client::set_capabilities(const ClientCapabilities& capabilities) {
    impl_->capabilities = capabilities;
}

const ClientCapabilities& Client::get_capabilities() const {
    return impl_->capabilities;
}

void Client::set_config(const ClientConfig& config) {
    impl_->config = config;
}

const ClientConfig& Client::get_config() const {
    return impl_->config;
}

void Client::set_notification_handler(std::function<void(const std::string&, const json::Value&)> handler) {
    impl_->notification_handler = std::move(handler);
}

void Client::set_error_handler(std::function<void(const Error&)> handler) {
    impl_->error_handler = std::move(handler);
}

Result<json::Value> Client::send_request(const json::Value& request) {
    std::promise<json::Value> promise;
    auto future = promise.get_future();

    {
        std::lock_guard<std::mutex> lock(impl_->message_mutex);
        auto request_id = impl_->generate_request_id();
        auto request_copy = request;
        request_copy["id"] = request_id;
        impl_->pending_requests[request_id] = std::move(promise);

        if (!impl_->transport->send_message(request_copy.serialize())) {
            impl_->pending_requests.erase(request_id);
            return Error::connection_error("Failed to send request");
        }
    }

    if (future.wait_for(std::chrono::seconds(impl_->config.timeout_seconds)) !=
        std::future_status::ready) {
        return Error::timeout("Request timed out");
    }

    auto response = future.get();
    if (response.is_object() && response.contains("error")) {
        auto error = response["error"];
        auto code = error.get("code", 0).as_int();
        auto message = error.get("message", "").as_string();
        return Error::protocol_error(message, static_cast<int>(code));
    }

    return response;
}

void Client::Impl::handle_message(const json::Value& message) {
    if (!message.is_object()) {
        return;
    }

    if (message.contains("id") && message.contains("result")) {
        handle_response(message);
    } else if (message.contains("method") && message.contains("params")) {
        auto method = message["method"].as_string();
        auto params = message["params"];

        if (message.contains("id")) {
            handle_request(message);
        } else {
            handle_notification(method, params);
        }
    }
}

void Client::Impl::handle_request(const json::Value& request) {
    // Handle server requests (ping, sampling, etc.)
    auto method = request["method"].as_string();
    auto id = request["id"].as_string();

    if (method == "ping") {
        auto response = json::Value::object({
            {"jsonrpc", "2.0"},
            {"id", id},
            {"result", json::Value::object({})}
        });

        transport->send_message(response.serialize());
    }
    // Handle other server requests as needed
}

void Client::Impl::handle_response(const json::Value& response) {
    auto id = response["id"].as_string();

    std::lock_guard<std::mutex> lock(message_mutex);
    auto it = pending_requests.find(id);
    if (it != pending_requests.end()) {
        it->second.set_value(response);
        pending_requests.erase(it);
    }
}

void Client::Impl::handle_notification(const std::string& method, const json::Value& params) {
    if (notification_handler) {
        notification_handler(method, params);
    }

    // Handle specific notifications
    if (method == "notifications/message") {
        // Handle resource updates, tool changes, etc.
    }
}

void Client::Impl::receiver_loop() {
    while (running) {
        try {
            auto message = transport->receive_message();
            if (!message) {
                break;
            }

            auto parsed = json::Value::parse(*message);
            if (!parsed) {
                continue;
            }

            handle_message(*parsed);
        } catch (const std::exception& e) {
            if (error_handler) {
                error_handler(Error::internal_error(std::string("Receiver error: ") + e.what()));
            }
            break;
        }
    }
}

// Factory functions
std::shared_ptr<Client> create_stdio_client(const ClientConfig& config) {
    auto transport = std::make_unique<StdioTransport>();
    return std::make_shared<Client>(std::move(transport), config);
}

std::shared_ptr<Client> create_http_client(const std::string& url, const ClientConfig& config) {
    auto transport = std::make_unique<HttpTransport>(url);
    return std::make_shared<Client>(std::move(transport), config);
}

std::shared_ptr<Client> create_websocket_client(const std::string& url, const ClientConfig& config) {
    auto transport = std::make_unique<WebSocketTransport>(url);
    return std::make_shared<Client>(std::move(transport), config);
}

std::shared_ptr<Client> create_sse_client(const std::string& url, const ClientConfig& config) {
    auto transport = std::make_unique<SseTransport>(url);
    return std::make_shared<Client>(std::move(transport), config);
}

} // namespace client
} // namespace mcp